package cn.edu.njust.nio;

import cn.edu.njust.EchoServer;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;

@Slf4j
public class NioEchoServer implements EchoServer {

    private ServerSocketChannel serverChannel;

    private Selector selector;

    private static String recv(SocketChannel socketChannel) {
        ByteArrayOutputStream bufferStream = new ByteArrayOutputStream();
        ByteBuffer buffer = ByteBuffer.allocate(1024);
        int len;

        try {
            // 这里注意条件写成 > 0 而不是 != -1
            while (((len = socketChannel.read(buffer)) > 0)) {
                // 记得翻转一下
                buffer.flip();
                bufferStream.write(buffer.array(), 0, buffer.limit());
            }

            return bufferStream.toString();
        } catch (IOException e) {
            e.printStackTrace();
        }

        return null;
    }

    private void send(SocketChannel socketChannel, String message) {

        try {
            socketChannel.write(ByteBuffer.wrap(message.getBytes()));
            // 发完消息记得关闭输出，否则会导致客户端read一直阻塞！(这里不采用拆包机制)
            socketChannel.shutdownOutput();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void listen(int port) {
        try {
            this.serverChannel = ServerSocketChannel.open();
            // 配置为非阻塞模式
            this.serverChannel.configureBlocking(false);
            this.serverChannel.bind(new InetSocketAddress(port));

            this.selector = Selector.open();
            // 注册客户端连接事件
            this.serverChannel.register(selector, SelectionKey.OP_ACCEPT);
        } catch (IOException e) {
            log.error("打开服务端通道失败");
            e.printStackTrace();
        }

    }

    @Override
    public void serve() {
        while (true) {
            try {
                this.selector.select();

                Iterator<SelectionKey> iterator = this.selector.selectedKeys().iterator();

                while (iterator.hasNext()) {
                    SelectionKey keySelected = iterator.next();

                    if (keySelected.isAcceptable()) {
                        // 连接请求
                        SocketChannel clientChannel = this.serverChannel.accept();
                        log.info("接收到来自客户端{}的请求", clientChannel.getRemoteAddress());
                        clientChannel.configureBlocking(false);
                        // 监听读消息事件
                        clientChannel.register(this.selector, SelectionKey.OP_READ);
                    } else if (keySelected.isReadable()) {
                        // 读客户端发来的消息
                        String message = recv((SocketChannel) keySelected.channel());
                        log.info("接收到来自客户端{}的消息{}", ((SocketChannel) keySelected.channel()).getRemoteAddress(), message);
                        send((SocketChannel) keySelected.channel(), message);
                        log.info("发送消息{}给{}", message, ((SocketChannel) keySelected.channel()).getRemoteAddress());
                        // 取消监听，否则会一直触发（Java NIO为水平触发）
                        keySelected.cancel();
                    }
                    // 记得remove
                    iterator.remove();
                }

            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}
